/*
 * Copyright 2024 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.djf.springaidemo.helper;

import jakarta.annotation.PreDestroy;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.model.Media;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.core.io.PathResource;
import org.springframework.stereotype.Component;
import org.springframework.util.MimeType;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static org.bytedeco.javacpp.Loader.deleteDirectory;

/**
 * FrameExtraHelper 是一个用于从视频中提取帧并将其缓存为图片文件的帮助类。
 * 它实现了 ApplicationRunner 接口，在应用启动时自动执行视频帧提取任务，
 * 并在应用关闭时清理临时文件。
 */
@Component
public final class FrameExtraHelper implements ApplicationRunner {

	// 私有构造函数防止外部实例化
	private FrameExtraHelper() {
	}

	/**
	 * 缓存提取出的图片路径列表，键为 "img"
	 */
	private static final Map<String, List<String>> IMAGE_CACHE = new ConcurrentHashMap<>();

	/**
	 * 视频文件路径（建议使用 classpath 或配置方式读取）
	 */
	private static final File videoUrl = new File("src/main/resources/multimodel/video.mp4");

	/**
	 * 提取后的帧图片保存路径（建议使用系统临时目录或可配置路径）
	 */
	private static final String framePath = "src/main/resources/multimodel/frame/";

	/**
	 * 日志记录器
	 */
	private static final Logger log = LoggerFactory.getLogger(FrameExtraHelper.class);

	/**
	 * 从视频中提取所有帧并保存为 PNG 图片到指定目录
	 */
	public static void getVideoPic() {
		List<String> strList = new ArrayList<>();
		File dir = new File(framePath);
		if (!dir.exists()) {
			dir.mkdirs();
		}

		try (
				FFmpegFrameGrabber ff = new FFmpegFrameGrabber(videoUrl.getPath());
				Java2DFrameConverter converter = new Java2DFrameConverter()
		) {
			ff.start();
			ff.setFormat("mp4");

			int length = ff.getLengthInFrames();

			Frame frame;
			for (int i = 1; i < length; i++) {
				frame = ff.grabFrame();
				if (frame.image == null) {
					continue;
				}
				BufferedImage image = converter.getBufferedImage(frame);
				String path = framePath + i + ".png";
				File picFile = new File(path);
				ImageIO.write(image, "png", picFile);
				strList.add(path);
			}
			IMAGE_CACHE.put("img", strList);
			ff.stop();
		} catch (Exception e) {
			log.error(e.getMessage());
		}
	}

	/**
	 * 应用启动时自动调用该方法，执行视频帧提取任务
	 *
	 * @param args 应用启动参数
	 * @throws Exception 如果发生异常
	 */
	@Override
	public void run(ApplicationArguments args) throws Exception {
		log.info("Starting to extract video frames");
		getVideoPic();
		log.info("Extracting video frames is complete");
	}

	/**
	 * 应用关闭前删除临时图片文件
	 */
	@PreDestroy
	public void destroy() {
		try {
			deleteDirectory(new File(framePath));
		} catch (IOException e) {
			log.error(e.getMessage());
		}
		log.info("Delete temporary files...");
	}

	/**
	 * 获取已提取的帧图片路径列表
	 *
	 * @return 帧图片路径列表
	 */
	public static List<String> getFrameList() {
		assert IMAGE_CACHE.get("img") != null;
		return IMAGE_CACHE.get("img");
	}

	/**
	 * 创建指定数量的 Media 对象列表，用于多模态模型输入
	 *
	 * @param numberOfImages 需要创建的 Media 对象数量
	 * @return Media 对象列表
	 */
	public static List<Media> createMediaList(int numberOfImages) {
		List<String> imgList = IMAGE_CACHE.get("img");
		int totalFrames = imgList == null ? 0 : imgList.size();
		int interval = Math.max(totalFrames / numberOfImages, 1);

		return IntStream.range(0, numberOfImages)
				.mapToObj(i -> imgList.get(i * interval))
				.map(image -> new Media(
						MimeType.valueOf("image/png"),
						new PathResource(image)
				))
				.collect(Collectors.toList());
	}
}
